/* Copyright (C) 2010 Luca Piccinelli
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
#include "imgops/skin/functions/skin_recognition.h"

using namespace cv;
using namespace std;

namespace NAMESPACE{

void skin_recognition_nbayes( const cv::Mat& src,
                              const NormalBayesModel& model,
                              cv::Mat& out_mask,
                              CVT_SPACE cvt_space,
                              int _class){

    CV_Assert( src.type()  == DataType<Vec3b>::type);

    Mat cvt_src;
    cvtColor(src, cvt_src, cvt_space);

    int x = 0, y = 0;
    switch (cvt_space){
        case YCRCB:
            x = 1, y = 2;
            break;
        case HSV:
            x = 0, y = 1;
            break;
        default:
        	return;
    }

    Mat responses, samples;
    cvt_src.reshape(1, src.rows * src.cols).colRange(x, y + 1).convertTo(samples, DataType<int>::type);

    model.predict(samples, responses);
    out_mask = responses.row(_class).reshape(1, src.rows);
}

void adaptive_skin_recognition_nbayes( const cv::Mat& src,
                                       const NormalBayesModel& model_off,
                                       const NormalBayesModel& model_on,
                                       cv::Mat& out_mask,
                                       double gain,
                                       CVT_SPACE cvt_space,
                                       int _class){

    if(model_on.is_trained()){
        CV_Assert( src.type()  == DataType<Vec3b>::type);

        Mat cvt_src;
        cvtColor(src, cvt_src, cvt_space);

        int x = 0, y = 0;
        switch (cvt_space){
            case YCRCB:
                x = 1, y = 2;
                break;
            case HSV:
                x = 0, y = 1;
                break;
        }

        const Discret2DNormalBayesModel* d_model_off = dynamic_cast<const Discret2DNormalBayesModel*>(&model_off);
        const Discret2DNormalBayesModel* d_model_on = dynamic_cast<const Discret2DNormalBayesModel*>(&model_on);

        Mat samples = cvt_src.reshape(1, src.rows * src.cols).colRange(x, y);
        Mat tmp_out(1, src.rows * src.cols, DataType<prob_t>::type);

        prob_t a_priori_evidence_prob_off, with_evidence_data_prob_off;
        prob_t a_priori_evidence_prob_on, with_evidence_data_prob_on;

        double gain2 = 1 - gain;
            
        a_priori_evidence_prob_off  = d_model_off->get_a_priori_evidence().at<prob_t>(0, 1);
        a_priori_evidence_prob_on   = d_model_on->get_a_priori_evidence().at<prob_t>(0, 1);
        for(int i = 0; i < samples.rows; i++){
            uchar* r = samples.ptr(i);
            int sample[2] = { r[0], r[1] };

            with_evidence_data_prob_off = d_model_off->get_with_evidence_data_luts().at(1).at<prob_t>(sample);
            with_evidence_data_prob_on  = d_model_on->get_with_evidence_data_luts().at(1).at<prob_t>(sample);

            // bayes classifier rule
            tmp_out.at<prob_t>(i) = ( (gain  * a_priori_evidence_prob_off * with_evidence_data_prob_off) / d_model_off->get_a_priori_data_lut().at<prob_t>(sample) ) +
                                    ( (gain2 * a_priori_evidence_prob_on  * with_evidence_data_prob_on)  / d_model_on->get_a_priori_data_lut().at<prob_t>(sample)  ) ;

        }
        out_mask = tmp_out.reshape(1, src.rows);
        
    }else{
        skin_recognition_nbayes(src, model_off, out_mask, cvt_space, _class);
    }
}

}
